import glob
import os

import torch
import tqdm
import time
import math
import copy

from torch.nn.utils import clip_grad_norm_
from pcdet.utils import common_utils, commu_utils
from pcdet.utils import self_training_utils
from pcdet.config import cfg

def train_one_epoch(model, optimizer, train_loader_1, train_loader_2, model_func, lr_scheduler, accumulated_iter, optim_cfg,
                    rank, tbar, total_it_each_epoch, dataloader_iter_1, dataloader_iter_2, tb_log=None, leave_pbar=False):
    if rank == 0:
        pbar = tqdm.tqdm(total=total_it_each_epoch, leave=leave_pbar, desc='train', dynamic_ncols=True)
        data_time = common_utils.AverageMeter()
        merge_time = common_utils.AverageMeter()
        batch_time = common_utils.AverageMeter()
        forward_time = common_utils.AverageMeter()

    for cur_it in range(total_it_each_epoch):
        end = time.time()
        try:
            batch_1 = next(dataloader_iter_1)
        except StopIteration:
            dataloader_iter_1 = iter(train_loader_1)
            batch_1 = next(dataloader_iter_1)
            #print('new iters')
        
        try:
            batch_2 = next(dataloader_iter_2)
        except StopIteration:
            dataloader_iter_2 = iter(train_loader_2)
            batch_2 = next(dataloader_iter_2)

        data_timer = time.time()
        cur_data_time = data_timer - end

        lr_scheduler.step(accumulated_iter)

        try:
            cur_lr = float(optimizer.lr)
        except:
            cur_lr = optimizer.param_groups[0]['lr']

        if tb_log is not None:
            tb_log.add_scalar('meta_data/learning_rate', cur_lr, accumulated_iter)

        model.train()
        optimizer.zero_grad()

        batch = common_utils.merge_two_batch_dict(batch_1, batch_2)
        merge_timer = time.time()
        cur_merge_time = merge_timer - data_timer

        loss, tb_dict, disp_dict = model_func(model, batch)

        forward_timer = time.time()
        cur_forward_time = forward_timer - data_timer

        loss.backward()
        clip_grad_norm_(model.parameters(), optim_cfg.GRAD_NORM_CLIP)
        optimizer.step()

        accumulated_iter += 1

        cur_batch_time = time.time() - end
        # average reduce
        avg_data_time = commu_utils.average_reduce_value(cur_data_time)
        avg_merge_time = commu_utils.average_reduce_value(cur_merge_time)
        avg_forward_time = commu_utils.average_reduce_value(cur_forward_time)
        avg_batch_time = commu_utils.average_reduce_value(cur_batch_time)

        # log to console and tensorboard
        if rank == 0:
            data_time.update(avg_data_time)
            merge_time.update(avg_merge_time)
            forward_time.update(avg_forward_time)
            batch_time.update(avg_batch_time)
            disp_dict.update({
                'loss': loss.item(), 'lr': cur_lr, 'd_time': f'{data_time.val:.2f}({data_time.avg:.2f})', 'm_time': f'{merge_time.val:.2f}({merge_time.avg:.2f})', 
                'f_time': f'{forward_time.val:.2f}({forward_time.avg:.2f})', 'b_time': f'{batch_time.val:.2f}({batch_time.avg:.2f})'
            })

            pbar.update()
            pbar.set_postfix(dict(total_it=accumulated_iter))
            tbar.set_postfix(disp_dict)
            tbar.refresh()

            if tb_log is not None:
                tb_log.add_scalar('train/loss', loss, accumulated_iter)
                tb_log.add_scalar('meta_data/learning_rate', cur_lr, accumulated_iter)
                for key, val in tb_dict.items():
                    tb_log.add_scalar('train/' + key, val, accumulated_iter)
    if rank == 0:
        pbar.close()
    return accumulated_iter


def train_model(model, optimizer, train_loader_1, train_loader_2, model_func, lr_scheduler, optim_cfg,
                start_epoch, total_epochs, start_iter, rank, tb_log, ckpt_save_dir, ps_label_dir,
                source_sampler=None, lr_warmup_scheduler=None, ckpt_save_interval=1,
                max_ckpt_save_num=50, merge_all_iters_to_one_epoch=False, logger=None):
    accumulated_iter = start_iter
    
    # Change the Data Augmentation: 
    # Cancel the GT-Sampling for the last 1/4 epoch for the training process
    
    # for resuming the ckpt: change the data_augmentator for dataset 1
    if cfg.DATA_CONFIG.get('PROG_AUG', None) and cfg.DATA_CONFIG.PROG_AUG.ENABLED and \
        (start_epoch == cfg.DATA_CONFIG.PROG_AUG.UPDATE_AUG_EPOCH):
        logger.info('**********************Starting the Fade GT-Sampling Operation**********************')
        train_loader_1.dataset.data_augmentor.re_prepare(
            augmentor_configs=cfg.DATA_CONFIG.PROG_AUG.D_CFG if cfg.DATA_CONFIG.PROG_AUG.get('D_CFG', None) else None)

    # for resuming the ckpt: change the data_augmentator for dataset 2
    if cfg.DATA_CONFIG_SRC_2.get('PROG_AUG', None) and cfg.DATA_CONFIG_SRC_2.PROG_AUG.ENABLED and \
        (start_epoch == cfg.DATA_CONFIG_SRC_2.PROG_AUG.UPDATE_AUG_EPOCH):
        logger.info('**********************Starting the Fade GT-Sampling Operation**********************')
        train_loader_2.dataset.data_augmentor.re_prepare(
            augmentor_configs=cfg.DATA_CONFIG_SRC_2.PROG_AUG.D_CFG if cfg.DATA_CONFIG_SRC_2.PROG_AUG.get('D_CFG', None) else None)

    with tqdm.trange(start_epoch, total_epochs, desc='epochs', dynamic_ncols=True, leave=(rank == 0)) as tbar:
        
        # generate the pseudo-labeling for merge the labeled A dataset and unlabeled B dataset
        if cfg.get('SELF_TRAIN', None):
            ps_pkl = self_training_utils.check_already_exsit_pseudo_label(ps_label_dir, start_epoch)
            if ps_pkl is not None:
                logger.info('==> Loading pseudo labels from {}'.format(ps_pkl))
        
        # it is wrong that total_it is equal to the sum of both train_loader_1 and train_loader_2
        # total_it_each_epoch = len(train_loader_1) + len(train_loader_2)
        total_it_each_epoch = len(train_loader_1) if len(train_loader_1) > len(train_loader_2) else len(train_loader_2)

        if merge_all_iters_to_one_epoch:
            raise NotImplementedError

        dataloader_iter_1 = iter(train_loader_1)
        dataloader_iter_2 = iter(train_loader_2)
        for cur_epoch in tbar:
            if source_sampler is not None:
                source_sampler.set_epoch(cur_epoch)

            # train one epoch
            if lr_warmup_scheduler is not None and cur_epoch < optim_cfg.WARMUP_EPOCH:
                cur_scheduler = lr_warmup_scheduler
            else:
                cur_scheduler = lr_scheduler
            
            if optim_cfg.get('GENERATE_PSEUDO_LABEL', None):
                train_loader_2.dataset.eval()
                logger.info('***********update pseudo label**********')
                self_training_utils.save_pseudo_label_epoch(
                    model, train_loader_2, rank,
                    leave_pbar=True, ps_label_dir=ps_label_dir, cur_epoch=cur_epoch
                )
                train_loader_2.dataset.train()

            # for resuming the ckpt: change the data_augmentator for dataset 1
            if cfg.DATA_CONFIG.get('PROG_AUG', None) and cfg.DATA_CONFIG.PROG_AUG.ENABLED and \
                (cur_epoch == cfg.DATA_CONFIG.PROG_AUG.UPDATE_AUG_EPOCH):
                logger.info('**********************Dataset ONE: Starting the Fade GT-Sampling Operation**********************')
                train_loader_1.dataset.data_augmentor.re_prepare(
                    augmentor_configs=cfg.DATA_CONFIG.PROG_AUG.D_CFG if cfg.DATA_CONFIG.PROG_AUG.get('D_CFG', None) else None)

            # for resuming the ckpt: change the data_augmentator for dataset 2
            if cfg.DATA_CONFIG_SRC_2.get('PROG_AUG', None) and cfg.DATA_CONFIG_SRC_2.PROG_AUG.ENABLED and \
                (cur_epoch == cfg.DATA_CONFIG_SRC_2.PROG_AUG.UPDATE_AUG_EPOCH):
                logger.info('**********************Dataset TWO: Starting the Fade GT-Sampling Operation**********************')
                train_loader_2.dataset.data_augmentor.re_prepare(
                    augmentor_configs=cfg.DATA_CONFIG_SRC_2.PROG_AUG.D_CFG if cfg.DATA_CONFIG_SRC_2.PROG_AUG.get('D_CFG', None) else None)

            accumulated_iter = train_one_epoch(
                model, optimizer, train_loader_1, train_loader_2, model_func,
                lr_scheduler=cur_scheduler,
                accumulated_iter=accumulated_iter, optim_cfg=optim_cfg,
                rank=rank, tbar=tbar, tb_log=tb_log,
                leave_pbar=(cur_epoch + 1 == total_epochs),
                total_it_each_epoch=total_it_each_epoch,
                dataloader_iter_1=dataloader_iter_1,
                dataloader_iter_2=dataloader_iter_2
            )

            # save trained model
            trained_epoch = cur_epoch + 1
            if trained_epoch % ckpt_save_interval == 0 and rank == 0:

                ckpt_list = glob.glob(str(ckpt_save_dir / 'checkpoint_epoch_*.pth'))
                ckpt_list.sort(key=os.path.getmtime)

                if ckpt_list.__len__() >= max_ckpt_save_num:
                    for cur_file_idx in range(0, len(ckpt_list) - max_ckpt_save_num + 1):
                        os.remove(ckpt_list[cur_file_idx])

                ckpt_name = ckpt_save_dir / ('checkpoint_epoch_%d' % trained_epoch)
                save_checkpoint(
                    checkpoint_state(model, optimizer, trained_epoch, accumulated_iter), filename=ckpt_name,
                )


def model_state_to_cpu(model_state):
    model_state_cpu = type(model_state)()  # ordered dict
    for key, val in model_state.items():
        model_state_cpu[key] = val.cpu()
    return model_state_cpu


def checkpoint_state(model=None, optimizer=None, epoch=None, it=None):
    optim_state = optimizer.state_dict() if optimizer is not None else None
    if model is not None:
        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
            model_state = model_state_to_cpu(model.module.state_dict())
        else:
            model_state = model.state_dict()
    else:
        model_state = None

    try:
        import pcdet
        version = 'pcdet+' + pcdet.__version__
    except:
        version = 'none'

    return {'epoch': epoch, 'it': it, 'model_state': model_state, 'optimizer_state': optim_state, 'version': version}


def save_checkpoint(state, filename='checkpoint'):
    if False and 'optimizer_state' in state:
        optimizer_state = state['optimizer_state']
        state.pop('optimizer_state', None)
        optimizer_filename = '{}_optim.pth'.format(filename)
        torch.save({'optimizer_state': optimizer_state}, optimizer_filename)

    filename = '{}.pth'.format(filename)
    torch.save(state, filename)